-- Lua 脚本实现令牌桶算法
-- key 为 Redis 中用于存储桶的键
-- rate 为令牌填充速率（每秒填充的令牌数）
-- capacity 为桶的容量
-- now 为当前时间戳
-- permits 为本次请求需要的令牌数

local key = KEYS[1]
local rate = tonumber(ARGV[1])
local capacity = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local permits = tonumber(ARGV[4])

local bucket = redis.call('hmget', key, 'lastRefillTime', 'tokens')
local lastRefillTime = tonumber(bucket[1])
local tokens = tonumber(bucket[2])

if lastRefillTime == nil then
    lastRefillTime = now
    tokens = capacity
end

-- 计算自上次填充以来经过的时间
local delta = math.max(0, now - lastRefillTime)
-- 计算应该填充的令牌数
local refillTokens = math.floor(delta * rate)
tokens = math.min(capacity, tokens + refillTokens)
lastRefillTime = now

local enoughTokens = false
if tokens >= permits then
    enoughTokens = true
    tokens = tokens - permits
end

-- 更新桶的状态
redis.call('hmset', key, 'lastRefillTime', lastRefillTime, 'tokens', tokens)

-- 设置过期时间防止无限增长
redis.call('expire', key, math.ceil(capacity/rate)*2)

if enoughTokens then
    return 1
else
    return 0
end